# ~*~ coding: utf-8 ~*~
import paramiko


class SSHConnection:
    connections = {}

    @staticmethod
    def make_key(user, asset, system_user):
        key = "{}_{}_{}".format(user.id, asset.id, system_user.id)
        return key

    @classmethod
    def new_connection_from_cache(cls, user, asset, system_user):
        key = cls.make_key(user, asset, system_user)
        connection = cls.connections.get(key)
        if not connection:
            return None
        if not connection.is_active:
            cls.connections.pop(key, None)
            return None
        connection.ref += 1
        return connection

    @classmethod
    def set_connection_to_cache(cls, conn):
        key = cls.make_key(conn.user, conn.asset, conn.system_user)
        cls.connections[key] = conn

    @classmethod
    def new_connection(cls, user, asset, system_user):
        connection = cls.new_connection_from_cache(user, asset, system_user)

        if connection:
            print("Reuse connection: {}->{}@{}".format(
                user.username, asset.manage_ip, system_user.username)
            )
            return connection
        # 缓存中没有就创建connection
        connection = cls(user, asset, system_user)
        connection.connect()
        if connection.is_active:
            cls.set_connection_to_cache(connection)
        return connection

    @classmethod
    def remove_ssh_connection(cls, conn):
        key = "{}_{}_{}".format(conn.user.id, conn.asset.id, conn.system_user.id)
        cls.connections.pop(key, None)

    def __init__(self, user, asset, system_user):
        self.user = user
        self.asset = asset
        self.system_user = system_user
        self.sock = None
        self.transport = None
        self.ref = 1

    def connect(self):
        ssh = paramiko.SSHClient()
        ssh.load_system_host_keys()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(self.asset.manage_ip, self.asset.port, self.system_user.username, self.system_user.password, timeout=5)
        transport = ssh.get_transport()
        transport.set_keepalive(20)
        self.transport = transport
        self.sock = ssh

    def reconnect_if_need(self):
        if not self.is_active:
            self.connect()

        if self.is_active:
            return True
        return False

    def get_transport(self):
        if self.reconnect_if_need():
            return self.transport
        return None

    def get_channel(self, term="xterm", width=80, height=24):
        if self.reconnect_if_need():
            chan = self.sock.invoke_shell(term, width=width, height=height)
            return chan
        else:
            return None

    def get_sftp(self):
        if self.reconnect_if_need():
            return self.sock.open_sftp()
        else:
            return None

    @property
    def is_active(self):
        return self.transport and self.transport.is_active()

    def close(self):
        if self.ref > 1:
            self.ref -= 1
            msg = "Connection ref -1: {}->{}@{}. {}".format(
                self.user.username, self.asset.hostname,
                self.system_user.username, self.ref
            )
            print(msg)
            return
        self.__class__.remove_ssh_connection(self)
        try:
            self.sock.close()
        except Exception as e:
            print("Close connection error: ", e)

        msg = "Close connection: {}->{}@{}. Total connections live: {}".format(
            self.user.username, self.asset.manage_ip,
            self.system_user.username, len(self.connections)
        )
        print(msg)

